# _*_ coding: utf-8 _*_
# Copyright MiqroEra, Inc 2022. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 (the "License")
# Python Script for LAMMPS Simulation
# Author: MiqroEra
# On working
from io import TextIOWrapper
from typing import List, Tuple, Optional


class Minimize():
    """
    The class is to minimize the energy of the system by adjusting atom coordinates.
    It may be used to relax the system after an equilibration run or to relax the system.
    It may provide a minimized local potential energy.
    """
    def __init__(self, style='cg', energy_tolerance=10**(-8), force_tolerance=10**(-10), max_iteration=10**7, max_evaluation=10**7) -> None:
        """
        The constructor of the class.
        :param style: The minimization style. The default is 'cg'.
        :param energy_tolerance: The energy tolerance. The default is 10**(-8).
        :param force_tolerance: The force tolerance. The default is 10**(-10).
        :param max_iteration: The maximum number of iterations. The default is 10**7.
        :param max_evaluation: The maximum number of force evaluations. The default is 10**7.
        """
        self._style = style
        self._energy_tolerance = energy_tolerance
        self._force_tolerance = force_tolerance
        self._max_iteration = max_iteration
        self._max_evaluation = max_evaluation

    def generate_minimize(self, file: TextIOWrapper) -> None:
        """
        The function is to generate the minimize command.
        :param file: The file to write the minimize command.
        :return: None
        """
        file.write('# Minimization Energy and Force.\n')
        file.write(f'{"minimize_style":<20} {self._style}\n')
        file.write(f'{"minimize":<20} {self._energy_tolerance} {self._force_tolerance} {self._max_iteration} {self._max_evaluation}\n')


class Equilibration():
    """
    The class is to equilibrate the system of amorphous polymers, refering: Theoretical Chemistry Accounts 132(3) (2009) 1–37.
    """

    def __init__(self, tet=300, tep=1, tmax=600, pmax=50000, 
                 tdamp="$(100.0*dt)", pdamp="$(100.0*dt)", 
                 nve_l_start=True, dump_fname="qeuil.lammpstrj",
                 dump_every=10000, dump_image=True, reset_timestp=True) -> None:
        """
        The constructor of the class.
        :param tet: The equilibration temperature. The default is 300.
        :param tep: The equilibration pressure. The default is 1.
        :param tmax: The maximum temperature. The default is 600.
        :param pmax: The maximum pressure. The default is 50000.
        :param tdamp: The temperature damping. The default is "$(100.0*dt)".
        :param pdamp: The pressure damping. The default is "$(100.0*dt)".
        :param nve_l_start: The flag of NVE-L start. The default is True.
        :param dump_fname: The dump file name. The default is "qeuil.lammpstrj".
        :param dump_every: The dump frequency. The default is 10000.
        :param dump_image: The flag of dump image. The default is True.
        :param reset_timestp: The flag of reset timestep. The default is True.
        """
        self._tet = tet
        self._tep = tep
        self._tmax = tmax
        self._pmax = pmax
        self._tdamp = tdamp
        self._pdamp = pdamp
        self._nve_l_start = nve_l_start

        duration = 0
        for i in self._eq_steps:
            duration += i[1]

        super().__init__(duration=duration, dump_fname=dump_fname, dump_every=dump_every, dump_image=dump_image, reset_timestp=reset_timestp)

    @property
    def _eq_steps(self):
        return [
            ['nvt', 50000, self._Tmax],
            ['nvt', 50000, self._Teq],
            ['npt', 50000, self._Teq, 0.02 * self._Pmax],
            ['nvt', 50000, self._Tmax],
            ['nvt', 100000, self._Teq],
            ['npt', 50000, self._Teq, 0.6 * self._Pmax],
            ['nvt', 50000, self._Tmax],
            ['nvt', 100000, self._Teq],
            ['npt', 50000, self._Teq, self._Pmax],
            ['nvt', 50000, self._Tmax],
            ['nvt', 100000, self._Teq],
            ['npt', 5000, self._Teq, 0.5 * self._Pmax],
            ['nvt', 5000, self._Tmax],
            ['nvt', 10000, self._Teq],
            ['npt', 5000, self._Teq, 0.1 * self._Pmax],
            ['nvt', 5000, self._Tmax],
            ['nvt', 10000, self._Teq],
            ['npt', 5000, self._Teq, 0.01 * self._Pmax],
            ['nvt', 5000, self._Tmax],
            ['nvt', 10000, self._Teq],
            ['npt', 800000, self._Teq, self._Peq],
        ]
    
    def generate_equilibration(self, file: TextIOWrapper) -> None:
        """
        The function is to generate the equilibration command.
        :param file: The file to write the equilibration command.
        :return: None
        """
        if self._nve_l_start:
            file.write('# NVE-L start.\n')
            file.write(f'{"fix":<20} fLANGEVIN all langevin '
                       f'{self._tmax} {self._tmax} {self._tdamp} \n')
            file.write(f'{"fix":<20} fNVELIMIT all nve/limit 0.1\n')
            file.write(f'{"fix":<20} fMOM all momentum 100 linear 1 1 1 angular\n')
            file.write(f'{"run":<20} 10000\n')
            file.write('\n')
            file.write(f'{"unfix":<20} fLANGEVIN\n')
            file.write(f'{"unfix":<20} fNVELIMIT\n')
            file.write(f'{"unfix":<20} fMOM\n')
            file.write('\n')
            file.write(f'{"reset_timestep":<15} 0\n')
            file.write('\n')

        for n, i enumerate(self._eq_steps):
            if i[0] == 'nvt':
                file.write(f'# NVT step {n}.\n')
                file.write(f'{"fix":<20} fLANGEVIN all langevin '
                           f'{i[2]} {i[2]} {self._tdamp} \n')
                file.write(f'{"run":<20} {i[1]}\n')
                file.write(f'{"unfix":<20} fLANGEVIN\n')
                file.write('\n')
            elif i[0] == 'npt':
                file.write(f'# NPT step {n}.\n')
                file.write(f'{"fix":<20} fLANGEVIN all langevin '
                           f'{i[2]} {i[2]} {self._tdamp} \n')
                file.write(f'{"fix":<20} fBAROSTAT all npt temp {i[2]} {i[2]} {self._tdamp} '
                           f'iso {i[3]} {i[3]} {self._pdamp} \n')
                file.write(f'{"run":<20} {i[1]}\n')
                file.write(f'{"unfix":<20} fLANGEVIN\n')
                file.write(f'{"unfix":<20} fBAROSTAT\n')
                file.write('\n')
            else:
                raise ValueError('The type of the equilibration step is not supported.')
        
        
